#include <opencv2/opencv.hpp>
#include <iostream>
#include <string>
#include <filesystem>
#include <pcl/visualization/pcl_visualizer.h>
#include <pcl/io/pcd_io.h>
#include <pcl/point_types.h>
#include <pcl/filters/passthrough.h>
#include <pcl/filters/statistical_outlier_removal.h>
#include <pcl/filters/voxel_grid.h>
#include <pcl/filters/median_filter.h>

#include "yoloInterface.h"

#define DEBUG

YOLOV8 *m_yolo;                                                  /// yolo检测模型
std::string m_YOLOmodelPath;                                     /// yolo模型权重文件
std::string m_YOLOClassName;                                     /// yolo检测的分类文件
double ppx = 314.973, ppy = 244.177, fx = 593.924, fy = 589.963; /// 相机内参

void yolo_init(std::string onnx_path, std::string classname_path)
{
    Config config;
    config.confThreshold = 0.2;
    config.nmsThreshold = 0.5;
    config.scoreThreshold = 0.5;
    config.inpWidth = 640;
    config.inpHeight = 480;
    config.onnx_path = onnx_path;

    m_yolo = new YOLOV8(config, m_YOLOClassName);
}

int main()
{
    // 1 获取当前项目路径
    std::string path = std::filesystem::current_path().parent_path().string();
    std::string rgb_path = path + "/images/color-0.21.png";
    std::string depth_path = path + "/images/depth-0.21.png";

#ifdef DEBUG
    std::cout << "current path is : " << path << std::endl;
#endif

    // 2 读取图片
    cv::Mat rgb_image = cv::imread(rgb_path, cv::IMREAD_COLOR);
    cv::Mat depth_image = cv::imread(depth_path, cv::IMREAD_UNCHANGED);

    // 3 初始化yolo模型
    m_YOLOmodelPath = path + "/config/best.onnx";
    m_YOLOClassName = path + "/config/classes.txt";
    yolo_init(m_YOLOmodelPath, m_YOLOClassName);

    // 4 yolo模型检测
    std::vector<outputDetection> m_outDetections; /// YOLOV8 getDetection接口输出的检测结果
    m_yolo->detect(rgb_image);                    // 检测
    m_outDetections.clear();
    m_outDetections = m_yolo->getDetection(); // 获取检测到的ROI区域

    // #ifdef DEBUG
    //     std::cout << "outDetections.size() = " << m_outDetections.size() << std::endl;
    //     cv::imshow("outDetections", rgb_image);
    //     cv::imshow("depth_image", depth_image);
    //     cv::waitKey(0);
    // #endif

    // 5 根据检测结果构造部分点云
    pcl::PointCloud<pcl::PointXYZRGB>::Ptr cloud(new pcl::PointCloud<pcl::PointXYZRGB>);

    for (int i = 0; i < m_outDetections.size(); i++)
    {
        cv::Rect cur_box = m_outDetections[i].box;
        std::cout << "识别到的对象的边框 = " << cur_box << std::endl;
        for (int m = cur_box.y; m < (cur_box.y + cur_box.height); m++)
        {
            for (int n = cur_box.x; n < (cur_box.x + cur_box.width); n++)
            {
                ushort d = depth_image.at<ushort>(m, n);
                if (d == 0)
                {
                    continue;
                }

                pcl::PointXYZRGB point;
                point.z = double(d) / 1000;
                point.x = (n - ppx) * point.z / fx;
                point.y = (m - ppy) * point.z / fy;

                // 计算颜色值
                cv::Vec3b bgr = rgb_image.at<cv::Vec3b>(m, n);
                point.r = bgr[2];
                point.g = bgr[1];
                point.b = bgr[0];

                cloud->push_back(point);
            }
        }
    }

    // #ifdef DEBUG
    //     pcl::visualization::PCLVisualizer viewer("Point Cloud Viewer");
    //     viewer.setBackgroundColor(10, 10, 10); // rgb
    //     viewer.addPointCloud(cloud, "cloud viewer");
    //     viewer.spin();
    // #endif

    // 6 点云处理
    // 6.1 直通滤波
    pcl::PassThrough<pcl::PointXYZRGB> pass; // 创建对象
    pass.setInputCloud(cloud);               // 设置输入点云
    pass.setFilterFieldName("z");            // 设置过滤字段，这里对z轴上的点云进行过滤
    pass.setFilterLimits(0.0, 2.5);          // 设置过滤范围，这里选择过滤掉0~2.5范围外的点云
    pass.filter(*cloud);                     // 执行滤波，结果保存在cloud_filter中
    // 6.2 统计滤波
    pcl::StatisticalOutlierRemoval<pcl::PointXYZRGB> sor; // 创建对象
    sor.setInputCloud(cloud);                             // 设置输入点云
    sor.setMeanK(50);                                     // 设置统计时考虑查询点邻近点数
    sor.setStddevMulThresh(1);                            // 设置为1代表：如果一个点的距离超过平均距离一个标准差以上，则会被当做离群点去除
    sor.filter(*cloud);                                   // 执行滤波，并将结果保存在cloud_filter中
    // 6.3 栅格滤波
    pcl::VoxelGrid<pcl::PointXYZRGB> vg; // 创建voxelGrid对象
    vg.setInputCloud(cloud);             // 设置输入点云
    vg.setLeafSize(0.01f, 0.01f, 0.01f); // 设置体素大小，单位是m，这里设置成了1cm的立方体
    vg.filter(*cloud);                   // 执行滤波，结果保存在cloud_filter中
    // 6.4 中值滤波
    pcl::MedianFilter<pcl::PointXYZRGB> mf;
    mf.setInputCloud(cloud);
    mf.setWindowSize(5);           // 过滤器窗口大小(窗口越大越平滑、但可能过度平滑)
    mf.setMaxAllowedMovement(0.05f); // 最大允许移动量（如果某点在滤波后的位置与原始位置超过了这个阈值，则被视为异常点不予滤波）
    mf.filter(*cloud);

#ifdef DEBUG
    pcl::visualization::PCLVisualizer filtered_viewer("after filter viewer");
    filtered_viewer.setBackgroundColor(10, 10, 10); // rgb
    filtered_viewer.addPointCloud(cloud, "cloud viewer");
    filtered_viewer.spin();
#endif

    return 0;
}
